import numpy as np
import pandas as pd
import matplotlib.pyplot as plt


# 计算样本与聚类中心的距离, 返回离簇中心最近的类别,返回是当前的样本数据属于那一个簇中心的id或者索引
def distance(sample, center_points):
    # 这里用差的平方来表示距离
    d = np.power(sample - center_points, 2).sum(axis=1)
    cls = d.argmin()
    return cls


# 对当前的分类子集进行可视化展示
def draw(clusters, step):
    color = ['red', 'blue', 'green']
    marker = ['+', 'x', '_']
    plt.figure(figsize=(5, 3.8))
    plt.title('第{}步'.format(step))
    plt.xlabel('密度', loc='center')
    plt.ylabel('含糖量', loc='center')
    # 用颜色区分k个簇的数据样本
    for i, cluster in enumerate(clusters):
        cluster = np.array(cluster)
        plt.scatter(cluster[:, 0], cluster[:, 1], c=color[i], marker=marker[i], s=150)
    plt.show()


# 根据输入的样本集与划分的簇数，分别返回k个簇样本,返回是每个簇的簇类中心
def k_means(samples, k):
    data_number = len(samples)
    centers_flag = np.zeros((k,))

    # 随机在数据中选择k个聚类中心
    center_points = samples[np.random.choice(data_number, k, replace=False)]

    step = 0
    while True:
        # 计算每个样本距离簇中心的距离, 然后分到距离最短的簇中心中
        clusters = [[] for i in range(k)]
        for sample in samples:
            ci = distance(sample, center_points)
            clusters[ci].append(sample)

        # 可视化当前的聚类结构
        draw(clusters, step)

        # 分完簇之后更新每个簇的中心点, 得到了簇中心继续进行下一步的聚类
        for i, sub_clusters in enumerate(clusters):
            new_center = np.array(sub_clusters).mean(axis=0)
            # 如果数值有变化则更新, 如果没有变化则设置标志位为1，当所有的标志位为1则退出循环
            if (center_points[i] != new_center).all():
                center_points[i] = new_center
            else:
                centers_flag[i] = 1

        step += 1
        center_points_str = ''
        for center_point in center_points:
            center_points_str += f'({center_point[0]:.6f}, {center_point[1]:.6f})\n'
        print(f'第{step}步 簇的中心点\n{center_points_str}')
        if centers_flag.all():
            break

    return center_points


# 根据簇类中心对簇进行分类，获取最后的分类结果
def split_data(samples, center_points):
    # 根据中心样本得知簇数
    k = len(center_points)
    clusters = [[] for i in range(k)]
    for sample in samples:
        ci = distance(sample, center_points)
        clusters[ci].append(sample)

    return clusters


def main():
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False

    np.random.seed(50)
    # 导入数据
    data = pd.read_csv(r'./dataset/watermelon-kmeans.txt')
    samples = data[['密度', '含糖率']].values

    center_points = k_means(samples=samples, k=3)
    clusters = split_data(samples=samples, center_points=center_points)
    for i, cluster in enumerate(clusters, 1):
        cluster_str = ''
        for point in cluster:
            cluster_str += f'({point[0]}, {point[1]}) '
        print(f'第{i}簇 {cluster_str}')


if __name__ == '__main__':
    main()
    